import numpy as np
import os
import proplot as plt
import cmaps

label_dir = r'E:\akesu_radar\cheng\20250701_label'
radar_dir = r'E:\akesu_radar\cheng\20250701'
label = '202'

labelfiles=sorted([os.path.join(label_dir,file) for file in os.listdir(label_dir) if file.endswith('.npy')])
radarfiles=sorted([os.path.join(radar_dir,file) for file in os.listdir(radar_dir) if file.endswith('.npz')])
print(len(labelfiles))
print(len(radarfiles))
nt=len(labelfiles)
picdir='./'

for it in range(nt):
    timestr=os.path.basename(labelfiles[it]).split('_')[0]
    print(timestr)
    lb=np.load(labelfiles[it])
    idx=np.where(lb==int(label))
    nn=idx[0].shape[0]
    print(nn)
    if nn==0:
        continue

    rd=np.load(radarfiles[it])
    dbz=rd['CR']
    lon=rd['lon']
    lat=rd['lat']

    tmp=np.where(lb==int(label),1,0)
    dbz=np.where(tmp==1,dbz,np.nan)

    fig,ax=plt.subplots(figsize=(6,6))
    h=ax.contourf(lon,lat,dbz,levels=np.arange(0,71,5),cmap=cmaps.radar)
    ax.format(title='dbz')
    ax.colorbar(h)
    picname=os.path.join(picdir,f'{timestr}.png')
    fig.canvas.print_figure(picname,dpi=500)
    


    
    